import pandas as pd
import statsmodels.formula.api as smf
from statsmodels.iolib.summary2 import summary_col

# read in file 'table_3_data.csv' which should be in the same directory as this code
tb3 = pd.read_csv('table_3_pseudo_data.csv')

# run regressions from table 3
dfs = tb3   # all states
reg1 = smf.ols(formula="rating_diff ~  moodys_first + C(week_diff_8)  -1 ", data=dfs).fit(cov_type='cluster', cov_kwds={'groups': dfs[['week_diff_8']]})
reg2 = smf.ols(formula="rating_diff ~  moodys_first + C(week_diff_8) + C(year) -1 ", data=dfs).fit(cov_type='cluster', cov_kwds={'groups': dfs[['week_diff_8']]})

dfs = tb3[ tb3.TX==1]   # texas only
reg3 = smf.ols(formula="rating_diff ~  moodys_first + C(week_diff_6)  -1 ", data=dfs).fit(cov_type='cluster', cov_kwds={'groups': dfs[['week_diff_6']]})
reg4 = smf.ols(formula="rating_diff ~  moodys_first + C(week_diff_6) + C(year) -1 ", data=dfs).fit(cov_type='cluster', cov_kwds={'groups': dfs[['week_diff_6']]})

# create table
results = summary_col([ reg1, reg2, reg3, reg4],stars=True,float_format='%0.3f',
                  model_names=['1','2','3','4'],
                  info_dict={'N':lambda x: "{0:d}".format(int(x.nobs)),
                             'R2':lambda x: "{:.2f}".format(x.rsquared)},
                             regressor_order=[  'moodys_first',
                                                'C(week_diff_8)[1.0]',
                                                'C(week_diff_6)[1.0]',
                                                'C(week_diff_8)[2.0]',
                                                'C(week_diff_6)[2.0]',
                                                'C(week_diff_8)[3.0]',
                                                'C(week_diff_6)[3.0]',
                                                'C(week_diff_8)[4.0]',
                                                'C(week_diff_6)[4.0]',
                                                'C(week_diff_8)[5.0]',
                                                'C(week_diff_6)[5.0]',
                                                'C(week_diff_8)[6.0]',
                                                'C(week_diff_6)[6.0]'
                                                'C(week_diff_8)[7.0]',
                                                'C(week_diff_8)[8.0]',],
                             drop_omitted=True)

# print table
print(results)
